import pandas as pd
import os
from pathlib import Path

import pandas as pd
import numpy as np

def append_to_jsonl(file_path, new_data, group_name):
    """
    将新数据追加到现有的JSONL文件

    参数:
        file_path: 文件路径
        new_data: 新数据DataFrame
        group_name: 分组名称
    """
    try:
        # 读取现有数据
        if file_path.stat().st_size > 0:  # 文件非空
            existing_df = pd.read_json(file_path, lines=True)
            # 合并数据
            combined_df = pd.concat([existing_df, new_data], ignore_index=True)
        else:
            combined_df = new_data

        combined_df =combined_df.assign(
            timestamp=convert_timestamp_simple((pd.to_datetime(combined_df['timestamp']).astype('int64'))) - 8 * 3600
        )

        # 写入合并后的数据
        combined_df.to_json(file_path, orient='records', lines=True, force_ascii=False)
        print(f"✓ 追加: {file_path.name} (+{len(new_data)} 行, 总计: {len(combined_df)} 行)")

    except Exception as e:
        print(f"❌ 追加失败 {file_path.name}: {e}")
        # 失败时回退到覆盖模式
        new_data.to_json(file_path, orient='records', lines=True, force_ascii=False)
        print(f"✓ 回退到覆盖: {file_path.name}")


def convert_timestamp_simple(series):
    """
    简化版：主要处理13位和10位时间戳
    """
    # 转换为数值
    timestamps = pd.to_numeric(series, errors='coerce')

    # 统一转换为10位秒级时间戳
    def convert_ts(ts):
        if pd.isna(ts):
            return ts

        ts_int = int(ts)

        # 基于数值范围判断时间戳单位
        if ts_int >= 10 ** 16:  # 纳秒级（>= 10^16）
            return ts_int // 10 ** 9  # 纳秒转秒
        elif ts_int >= 10 ** 13:  # 微秒级（10^13-10^15）
            return ts_int // 10 ** 6  # 微秒转秒
        elif ts_int >= 10 ** 10:  # 毫秒级（10^10-10^12）
            return ts_int // 1000  # 毫秒转秒
        elif ts_int >= 10 ** 9:  # 秒级（10^9-10^10）
            return ts_int
        else:
           return ts_int  # 转换失败返回原值

    result = timestamps.apply(convert_ts)
    return result.astype('Int64')



def fast_grouped_processing(df, group_col, key_mapping, timestamp_col,code_list, output_dir='grouped_output',types="15min"):
    """
    快速分组处理DataFrame并写入多个JSONL文件

    参数:
        df: 输入DataFrame
        group_col: 分组列名
        key_mapping: 列重命名映射
        timestamp_col: 时间戳列名
        output_dir: 输出目录
    """

    # 创建输出目录
    # os.makedirs(output_dir, exist_ok=True)

    td = types.split("m")[0]

    # row["y"] = year
    # row["m"] = month
    # row["d"] = day
    # row["td"] = day
    # row["timestamp"] = timestamp
    # row["pc"] = 0.0
    # row["sf"] = 0

    # 预处理：重命名和时间提取
    processed_df = (df
    .rename(columns=key_mapping)
    .assign(
        y=lambda x: pd.to_datetime(x[timestamp_col]).dt.year,
        # m=lambda x: pd.to_datetime(x[timestamp_col]).dt.strftime('%Y%m'),
        # d=lambda x: pd.to_datetime(x[timestamp_col]).dt.strftime('%Y%m%d'),
        m=lambda x: (lambda dt: dt.dt.year * 100 + dt.dt.month)(pd.to_datetime(x[timestamp_col])),
        d=lambda x: (lambda dt: dt.dt.year * 10000 + dt.dt.month * 100 + dt.dt.day)(pd.to_datetime(x[timestamp_col])),
        td=str(td+""),
        pc =0.0,
        sf=0,
        # timestamp = (pd.to_datetime(df[timestamp_col]).astype('int64') // 10 ** 9).astype(int) - 8 * 3600
        timestamp = convert_timestamp_simple((pd.to_datetime(df[timestamp_col]).astype('int64'))) - 8 * 3600
    )
    )

    print("Value:", processed_df['td'].iloc[0])
    print("Type of element:", type(processed_df['td'].iloc[0]))
    print("Dtype:", processed_df['td'].dtype)


    # 按组处理并写入文件
    for group_name, group_data in processed_df.groupby(group_col):
        # 移除分组列（避免重复）
        group_data_without_key = group_data.drop(columns=[group_col,timestamp_col])

        #不在指定的股票列表里面 直接跳出
        if len(code_list)>0 and (group_name not in code_list):
            continue

        # 最简版本：直接在你的to_json调用前插入这行
        group_data_without_key = group_data_without_key[["t", "td", "timestamp", "y", "m", "d", "o", "h", "l", "c", "v", "a", "pc", "sf"]]

        # 确保输出目录存在
        output_path = Path(output_dir) / str(group_name)
        if not output_path.exists():
            output_path.mkdir(parents=True, exist_ok=True)



        etf_code =group_name.split(".")[0]
        # 写入JSONL文件
        output_file = output_path / f"{etf_code}_{td}.jsonl"

        if output_file.exists() and output_file.stat().st_size > 0:
            # 追加模式：读取现有数据，合并，然后写入
            append_to_jsonl(output_file, group_data_without_key, group_name)
        else:
            # 新文件：直接写入
            group_data_without_key.to_json(
                str(output_file),
                orient='records',
                lines=True,
                force_ascii=False
            )
        print(f"✓ 创建: {output_file.name} ({len(group_data)} 行)")


        print(f"✓ 已写入: {output_file} ({len(group_data)} 行)")

    print(f"\n✅ 处理完成！共处理 {len(processed_df)} 行数据")
    return len(processed_df.groupby(group_col))


# 使用示例
# df = pd.DataFrame({
#     'code': ['A001', 'A001', 'B002', 'B002', 'C003', 'C003'],
#     'user_id': [1, 2, 3, 4, 5, 6],
#     'user_name': ['Alice', 'Bob', 'Charlie', 'Diana', 'Eve', 'Frank'],
#     'create_time': [
#         '2023-01-15 10:30:00', '2023-02-20 14:45:00',
#         '2023-03-25 09:15:00', '2023-04-10 16:20:00',
#         '2023-05-15 11:30:00', '2023-06-20 13:45:00'
#     ]
# })
#
# key_mapping = {'user_id': 'id', 'user_name': 'name'}
# result = fast_grouped_processing(df, 'code', key_mapping, 'create_time')